{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5b2fe72-599d-46ed-8285-25e9e0d31ffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import pickle\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "\n",
    "from utils import *\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9142f7d-b27e-4b33-9e69-0b813215e073",
   "metadata": {},
   "outputs": [],
   "source": [
    "path_base = '../data/example_images'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27f30403-0150-4722-a140-274db3fb4f44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load gaze estimation results\n",
    "with open('../inference_results/vertex/ALL/test_0/inference/results_vertex_mode.pkl', 'rb') as f:\n",
    "    results = pickle.load(f)\n",
    "results = {k: v for k, v in results.items()}\n",
    "paths = list(results.keys())\n",
    "len(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a09a8522-efb4-4d41-aec7-85cd68c76ef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load preprocessing data\n",
    "with open('../output/preprocessing/data_face68.pkl', 'rb') as f:\n",
    "    data_pproc = pickle.load(f)\n",
    "data_pproc = {d['name']: d for d in data_pproc}\n",
    "len(data_pproc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07baab32-872f-4b1e-8bde-3448e9e5049c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Randomly loop through the results\n",
    "k = np.random.randint(len(paths))\n",
    "path = paths[k]\n",
    "\n",
    "res = results[path]\n",
    "path = '../' + path\n",
    "image = cv2.imread(path)\n",
    "# lms68 = data_pproc[os.path.relpath(path, path_base)]['face']['xyz68'].astype(np.float32)\n",
    "lms5 = data_pproc[os.path.relpath(path, path_base)]['face']['xy5'].astype(np.float32)\n",
    "eyes = {\n",
    "    'left' : (res['eyes']['left']['P'] @ eyel_template_homo.T).T,\n",
    "    'right': (res['eyes']['right']['P'] @ eyer_template_homo.T).T\n",
    "}\n",
    "gaze_vector = res['gaze_vec_combined']\n",
    "print(f\"Gaze direction: {gaze_vector}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfcb10b6-1736-41e7-8aa1-94aa60e081d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize gaze direction\n",
    "# image_gaze = draw_gaze_from_vector(image.copy(), lms68, gaze_vector, colour=[255, 0, 0])\n",
    "image_gaze = draw_gaze_from_vector(image.copy(), lms5, gaze_vector, colour=[255, 0, 0])\n",
    "plt.imshow(image_gaze[:, :, [2, 1, 0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb35015f-3e74-4f6b-a812-4424e03fd6eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize 3D eyes\n",
    "# image_eyes = draw_eyes(image.copy(), lms68, eyes, colour=[178, 255, 102])\n",
    "image_eyes = draw_eyes(image.copy(), lms5, eyes, colour=[178, 255, 102])\n",
    "plt.imshow(image_eyes[:, :, [2, 1, 0]])\n",
    "image.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
